倍增法, RMQ

LeetCode 2836. 在传球游戏中最大化函数值

给你一个长度为 n 下标从 0 开始的整数数组 receiver 和一个整数 k

总共有 n 名玩家,玩家 编号 互不相同,且为 [0, n - 1] 中的整数。这些玩家玩一个传球游戏,receiver[i] 表示编号为 i 的玩家会传球给编号为 receiver[i] 的玩家。玩家可以传球给自己,也就是说 receiver[i] 可能等于 i

你需要从 n 名玩家中选择一名玩家作为游戏开始时唯一手中有球的玩家,球会被传 恰好 k 次。

如果选择编号为 x 的玩家作为开始玩家,定义函数 f(x) 表示从编号为 x 的玩家开始,k 次传球内所有接触过球玩家的编号之 和 ,如果有玩家多次触球,则 累加多次 。换句话说, f(x) = x + receiver[x] + receiver[receiver[x]] + ... + receiver(k)[x]

你的任务时选择开始玩家 x ,目的是 最大化 f(x)

请你返回函数的 最大值 。

注意:receiver 可能含有重复元素。

示例 1:

传递次数 传球者编号 接球者编号 x + 所有接球者编号
 	 	 	2
1	2	1	3
2	1	0	3
3	0	2	5
4	2	1	6

输入:receiver = [2,0,1], k = 4
输出:6
解释:上表展示了从编号为 x = 2 开始的游戏过程。
从表中可知,f(2) 等于 6 。
6 是能得到最大的函数值。
所以输出为 6 。

示例 2:

传递次数 传球者编号 接球者编号 x + 所有接球者编号
 	 	 	4
1	4	3	7
2	3	2	9
3	2	1	10
 
输入:receiver = [1,1,1,2,3], k = 3
输出:10
解释:上表展示了从编号为 x = 4 开始的游戏过程。
从表中可知,f(4) 等于 10 。
10 是能得到最大的函数值。
所以输出为 10 。

提示:

  • 1 <= receiver.length == n <= 10^5
  • 0 <= receiver[i] <= n - 1
  • 1 <= k <= 10^10
cpp
// 倍增法 模板
typedef long long LL;
const int N = 1e5 + 10, M = 35;
int f[N][M];
LL val[N][M];
class Solution {
public:
    long long getMaxFunctionValue(vector<int>& a, long long k) {
        int n = a.size();
        for(int i = 0; i < n ; i ++ ) {
            f[i][0] = val[i][0] = a[i];
        }
        for(int j = 1; j < M ; j ++ ){
            for(int i = 0; i < n ; i ++ ){
                f[i][j] = f[f[i][j - 1]][j - 1];
                val[i][j] = val[f[i][j - 1]][j - 1] + val[i][j - 1];
            }
        }
        LL res = 0;
        for(int i = 0; i < n ; i ++ ){
            LL cur = i;
            int u = i;
            for(int j = M - 1; j >= 0 ; j -- ){
                if(k >> j & 1){
                    cur += val[u][j];
                    u = f[u][j];
                }
            }
            res = max(res, cur);
        }
        return res;
    }
};

python
'''
# pa[x][i] 表示 x 的第 2 ^ i 个祖先节点
pa[x][0] = parent[x]
pa[x][1] = pa[pa[x][0]][0]
pa[x][i + 1] = pa[pa[x][i]][i]
pa[x][i] = (p, s) 表示从 x 向上跳 2 ^ i 步能到达 p 点, 以及节点和 s 
'''
class Solution:
    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:
        n = len(receiver)
        m = 35
        pa = [[(p, p)] + [None] * m for p in receiver]
        for i in range(m):
            for x in range(n):
                p, s = pa[x][i] # 2 ^ i
                pp, ss = pa[p][i] # 2 ^ i
                pa[x][i + 1] = (pp, s + ss) # 2 ^ (i + 1)
        res = 0
        for i in range(n):
            x = sum = i # 从这个 x 节点向上跳, 这个节点编号为 i (因为答案是编号和)
            for j in range(m):
                if (k >> j) & 1: # 需要向上跳 2 ^ j 步
                    x, s = pa[x][j] # 2 ^ j
                    sum += s
            res = max(res, sum)
        return res

cpp
using i64 = long long;
const int N = 1e5 + 10, M = 36;

i64 pa[N][M], val[N][M];

class Solution {
public:
    long long getMaxFunctionValue(vector<int>& receiver, long long k) {
        int n = receiver.size();
        for (int i = 0; i < n ; i ++ ) {
            pa[i][0] = val[i][0] = receiver[i];
        }
        for (int j = 1; j <= 35 ; j ++ ) {
            for (int i = 0; i < n ; i ++ ) {
                i64 p = pa[i][j - 1], s = val[i][j - 1];
                pa[i][j] = pa[p][j - 1];
                val[i][j] = val[p][j - 1] + s;
            }
        }
        i64 res = 0;
        for (int i = 0; i < n ; i ++ ) {
            i64 p = i, sum = i;
            for (int j = 0; j <= 35 ; j ++ ) {
                if (k >> j & 1) {
                    sum += val[p][j];
                    p = pa[p][j];
                }
            }
            res = max(res, sum);
        }
        return res;
    }
};